import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from comet import download_model, load_from_checkpoint


model_path = download_model("Unbabel/unite-mup")
model = load_from_checkpoint(model_path)
data = [
    {
        "src": "这是个句子。",
        "mt": "This is a sentence.",
        # "ref": "It is a sentence."
        "ref": "a phrase. "
    },
    {
        "src": "这是个句子。",
        "mt": "this is phrase.",
        "ref": "It is a sentence."
    },
    {
        "src": "这是另一个句子。",
        "mt": "This is another sentence.",
        "ref": "It is another sentence."
    }
]
model_output = model.predict(data, batch_size=8, gpus=1)
print (model_output)

# Expected SRC score:
# [0.3474583327770233, 0.4492775797843933]
print (model_output.metadata.src_scores)

# Expected REF score:
# [0.9252626895904541, 0.899452269077301]
print (model_output.metadata.ref_scores)

# Expected UNIFIED score:
# [0.8758717179298401, 0.8294666409492493]
print (model_output.metadata.unified_scores)